import sys

import pandas as pd

from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import GridSearchCV


# print(pd.__version__)

class RF_predict:
    address = ''

    def __init__(self, address):
        self.address = address

    # 找到最优参数
    def get_best_params(self, X_train, y_train):
        parameters = {'n_estimators': [5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 20], 'max_depth': [2, 3, 4, 5, 6],
                      'min_samples_leaf': [5, 10, 20, 30, 40]}
        new_model = RandomForestRegressor(random_state=1)

        grid_search = GridSearchCV(new_model, parameters, cv=6, scoring='accuracy')
        grid_search.fit(X_train, y_train)
        a = grid_search.best_params_
        print(a)

    # 短期均线能量计算
    def make_predict_short(self):

        df = pd.read_excel(self.address)

        # 设置自变量和因变量
        # 自变量为10线、20线、排列金叉死叉、斜率、距离
        # 因变量为价格的变化量
        X = df[['MA10', 'MA20', 'XL10', 'XL20', 'S1020']]
        y = df[['endP']]

        # 设置训练集和测试集
        X_len = X.shape[0]
        split = int(X_len * 0.9)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        date = df[['everyDate']]
        dateList = date[split:]

        # 创建模型
        model = RandomForestRegressor(n_estimators=500, random_state=1, bootstrap=True, oob_score=False, n_jobs=1)
        model.fit(X_train, y_train)

        # 用测试集测试准确性
        y_pred = model.predict(X_test)
        a = pd.DataFrame()

        # 处理索引（因为测试集是一个切片，索引从切掉的地方计数，因此需要重置索引）
        y_test_index = list()
        for i in range(len(y_test)):
            y_test_index.append(i)
        y_test.index = y_test_index

        a['预测值'] = list(y_pred)
        a['实际值'] = y_test

        self.writeResult(destFileName + "__S", dateList, y_pred, y_test)

        # return a

    # 中期均线能量
    def make_predict_mid(self):

        df = pd.read_excel(self.address)

        # 设置自变量和因变量
        # 自变量为10线、20线、排列金叉死叉、斜率、距离
        # 因变量为价格的变化量
        X = df[['MA30', 'MA60', 'XL30', 'XL60', 'S3060']]
        y = df[['endP']]

        # 设置训练集和测试集
        X_len = X.shape[0]
        split = int(X_len * 0.9)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        date = df[['everyDate']]
        dateList = date[split:]

        # 创建模型
        model = RandomForestRegressor(n_estimators=500, random_state=1, bootstrap=True, oob_score=False, n_jobs=1)
        model.fit(X_train, y_train)

        # 用测试集测试准确性
        y_pred = model.predict(X_test)
        a = pd.DataFrame()

        # 处理索引（因为测试集是一个切片，索引从切掉的地方计数，因此需要重置索引）
        y_test_index = list()
        for i in range(len(y_test)):
            y_test_index.append(i)
        y_test.index = y_test_index

        a['预测值'] = list(y_pred)
        a['实际值'] = y_test

        # # 分析数据特征的重要性
        # features = X.columns
        # importances = model.feature_importances_
        # a['特征'] = features
        # a['特征重要性'] = importances

        self.writeResult(destFileName + "__M", dateList, y_pred, y_test)

        # print(a)
        # return a

    # 长期均线能量
    def make_predict_long(self):

        df = pd.read_excel(self.address)

        # 设置自变量和因变量
        # 自变量为10线、20线、排列金叉死叉、斜率、距离
        # 因变量为价格的变化量
        X = df[['MA120', 'MA250', 'XL120', 'XL250', 'S120250']]
        y = df[['endP']]

        # 设置训练集和测试集
        X_len = X.shape[0]
        split = int(X_len * 0.9)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        date = df[['everyDate']]
        dateList = date[split:]

        # 创建模型
        model = RandomForestRegressor(n_estimators=500, random_state=1, bootstrap=True, oob_score=False, n_jobs=1)
        model.fit(X_train, y_train)

        # 用测试集测试准确性
        y_pred = model.predict(X_test)
        a = pd.DataFrame()

        # 处理索引（因为测试集是一个切片，索引从切掉的地方计数，因此需要重置索引）
        y_test_index = list()
        for i in range(len(y_test)):
            y_test_index.append(i)
        y_test.index = y_test_index

        a['预测值'] = list(y_pred)
        a['实际值'] = y_test

        # # 分析数据特征的重要性
        # features = X.columns
        # importances = model.feature_importances_
        # a['特征'] = features
        # a['特征重要性'] = importances

        self.writeResult(destFileName + "__L", dateList, y_pred, y_test)

        # print(a)
        # return a

    # 短中期均线能量
    def make_predict_short_mid(self):

        df = pd.read_excel(self.address)

        # 设置自变量和因变量
        # 自变量为10线、20线、排列金叉死叉、斜率、距离
        # 因变量为价格的变化量
        X = df[['MA10', 'MA20', 'XL10', 'XL20', 'S1020', 'MA30', 'MA60', 'XL30', 'XL60', 'S3060']]
        y = df[['endP']]

        # 设置训练集和测试集
        X_len = X.shape[0]
        split = int(X_len * 0.9)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        date = df[['everyDate']]
        dateList = date[split:]

        # 创建模型
        model = RandomForestRegressor(n_estimators=500, random_state=1, bootstrap=True, oob_score=False, n_jobs=1)
        model.fit(X_train, y_train)

        # 用测试集测试准确性
        y_pred = model.predict(X_test)
        a = pd.DataFrame()

        # 处理索引（因为测试集是一个切片，索引从切掉的地方计数，因此需要重置索引）
        y_test_index = list()
        for i in range(len(y_test)):
            y_test_index.append(i)
        y_test.index = y_test_index

        a['预测值'] = list(y_pred)
        a['实际值'] = y_test

        # # 分析数据特征的重要性
        # features = X.columns
        # importances = model.feature_importances_
        # a['特征'] = features
        # a['特征重要性'] = importances

        self.writeResult(destFileName + "__SM", dateList, y_pred, y_test)

        # print(a)
        # return a

    # 中长期均线能量
    def make_predict_mid_long(self):

        df = pd.read_excel(self.address)

        # 设置自变量和因变量
        # 自变量为10线、20线、排列金叉死叉、斜率、距离
        # 因变量为价格的变化量
        X = df[['MA30', 'MA60', 'XL30', 'XL60', 'S3060', 'MA120', 'MA250', 'XL120', 'XL250', 'S120250']]
        y = df[['endP']]

        # 设置训练集和测试集
        X_len = X.shape[0]
        split = int(X_len * 0.9)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        date = df[['everyDate']]
        dateList = date[split:]

        # 创建模型
        model = RandomForestRegressor(n_estimators=500, random_state=1, bootstrap=True, oob_score=False, n_jobs=1)
        model.fit(X_train, y_train)

        # 用测试集测试准确性
        y_pred = model.predict(X_test)
        a = pd.DataFrame()

        # 处理索引（因为测试集是一个切片，索引从切掉的地方计数，因此需要重置索引）
        y_test_index = list()
        for i in range(len(y_test)):
            y_test_index.append(i)
        y_test.index = y_test_index

        a['预测值'] = list(y_pred)
        a['实际值'] = y_test

        # # 分析数据特征的重要性
        # features = X.columns
        # importances = model.feature_importances_
        # a['特征'] = features
        # a['特征重要性'] = importances

        self.writeResult(destFileName + "__ML", dateList, y_pred, y_test)

        # print(a)
        # return a

    # 三期均线能量
    def make_predict_short_mid_long(self):

        df = pd.read_excel(self.address)

        # 设置自变量和因变量
        # 自变量为10线、20线、排列金叉死叉、斜率、距离
        # 因变量为价格的变化量
        X = df[['MA10', 'MA20', 'XL10', 'XL20', 'S1020', 'MA30', 'MA60', 'XL30', 'XL60', 'S3060', 'MA120', 'MA250',
                'XL120', 'XL250', 'S120250']]
        y = df[['endP']]

        # 设置训练集和测试集
        X_len = X.shape[0]
        split = int(X_len * 0.9)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        date = df[['everyDate']]
        dateList = date[split:]

        # 创建模型
        model = RandomForestRegressor(n_estimators=500, random_state=1, bootstrap=True, oob_score=False, n_jobs=1)
        model.fit(X_train, y_train)

        # 用测试集测试准确性
        y_pred = model.predict(X_test)
        a = pd.DataFrame()

        # 处理索引（因为测试集是一个切片，索引从切掉的地方计数，因此需要重置索引）
        y_test_index = list()
        for i in range(len(y_test)):
            y_test_index.append(i)
        y_test.index = y_test_index

        a['预测值'] = list(y_pred)
        a['实际值'] = y_test

        # # 分析数据特征的重要性
        # features = X.columns
        # importances = model.feature_importances_
        # a['特征'] = features
        # a['特征重要性'] = importances

        self.writeResult(destFileName + "__SML", dateList, y_pred, y_test)

        # print(a)
        # return a

    def writeResult(self, destFileName, dateList, y_pred, y_test):
        with open(destFileName + ".txt", "w") as f:
            y_predList = list(y_pred)
            y_testList = y_test.values.tolist()
            dateList = dateList.values.tolist()
            for i in range(len(y_predList)):
                f.write(dateList[i][0] + "\t")
                f.write(str(round(y_predList[i], 2)) + "\t")
                f.write(str(y_testList[i][0]) + "\n")


fileName = sys.argv[1]
lastIndexOfName = fileName.rfind("/")
lastIndexOfPoint = fileName.rfind(".")
destFileName = fileName[lastIndexOfName + 1:lastIndexOfPoint]

# x = RF_predict('E:/桌面/毕业设计/股票数据/深b股数据/res/ST山航.xls')
x = RF_predict(fileName)
x.make_predict_short()
x.make_predict_mid()
x.make_predict_long()
x.make_predict_short_mid()
x.make_predict_mid_long()
x.make_predict_short_mid_long()
